using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using Unity.MLAgents;

public class EnvParaManager
{
    private EnvironmentParameters eps;

    // The envParaDict stores the current environment parameters' value.
    private Dictionary<EP, EnvPara> envParaDict;

    //Initialization of environment parameters
    //After the initialization of EnvParaManager, 
    //there is only one way to modify the value of element in envParaDict: EnvironmentSettingChannel
    //standard_fly version settings
    // public EnvParaManager()
    // {
    //     eps = Academy.Instance.EnvironmentParameters;

    //     envParaDict = new Dictionary<EP, EnvPara>();

    //     envParaDict.Add(EP.seed, new EnvPara(EnvParaType.INT, 5050));
    //     envParaDict.Add(EP.playgroundNumber, new EnvPara(EnvParaType.INT, 10));

    //     //*All the position below is localPosition, which being relative to the playground gameObject
    //     //Target position  
    //     envParaDict.Add(EP.targetPosX, new EnvPara(EnvParaType.RANDOM_UNIFORM, new List<float>{-20,20,0}));
    //     envParaDict.Add(EP.targetPosY, new EnvPara(EnvParaType.RANDOM_UNIFORM, new List<float>{10,30,20}));
    //     envParaDict.Add(EP.targetPosZ, new EnvPara(EnvParaType.FLOAT, 150f));

    //     //Agent's spawn position 
    //     envParaDict.Add(EP.spawnPosX, new EnvPara(EnvParaType.FLOAT, 0f));
    //     envParaDict.Add(EP.spawnPosY, new EnvPara(EnvParaType.FLOAT, 10f));
    //     envParaDict.Add(EP.spawnPosZ, new EnvPara(EnvParaType.FLOAT, 0f));
        
    //     //Set random rotation every time agent respawn
    //     envParaDict.Add(EP.enableRandomAirplaneSpawnRotation, new EnvPara(EnvParaType.BOOL, false));

    //     //Agents should only fly in a fixed cuboid area.
    //     //If agents fly out the area, the task failed.
    //     envParaDict.Add(EP.enableBoundaryCheck, new EnvPara(EnvParaType.BOOL, true));
    //     envParaDict.Add(EP.boundaryForward, new EnvPara(EnvParaType.FLOAT, 170f));
    //     envParaDict.Add(EP.boundaryBackward, new EnvPara(EnvParaType.FLOAT, -10f));
    //     envParaDict.Add(EP.boundaryUp, new EnvPara(EnvParaType.FLOAT, 40f));
    //     envParaDict.Add(EP.boundaryDown, new EnvPara(EnvParaType.FLOAT, 0f));
    //     envParaDict.Add(EP.boundaryLeft, new EnvPara(EnvParaType.FLOAT, -40f));
    //     envParaDict.Add(EP.boundaryRight, new EnvPara(EnvParaType.FLOAT, 40f));

    //     //If agents is too far from target, the task failed.
    //     envParaDict.Add(EP.maxDistanceToTarget, new EnvPara(EnvParaType.FLOAT, 1000f));

    //     //The number of agent per playground, only support 1 recently.
    //     envParaDict.Add(EP.agentNumberPerPlayground, new EnvPara(EnvParaType.INT, 1));

    //     //z coordinate of badSpheres. Also indicate the number of badSpheres.
    //     envParaDict.Add(EP.badSpheresPosZ, new EnvPara(EnvParaType.LIST_FLOAT, new List<float>{40,60,80,100}));


    //     //The x,y coordinates and scale of badSphere is random. 
    //     envParaDict.Add(EP.badSpherePosX, new EnvPara(EnvParaType.RANDOM_UNIFORM, new List<float>{-35,35,0}));
    //     envParaDict.Add(EP.badSpherePosY, new EnvPara(EnvParaType.RANDOM_UNIFORM, new List<float>{5,35,20}));
    //     envParaDict.Add(EP.badSphereScale, new EnvPara(EnvParaType.RANDOM_UNIFORM, new List<float>{15,15,15}));

    //     //The scale of the distance reward
    //     //distance reward(every action) = Tanh(distanceToTarget * 0.001f) * scale
    //     envParaDict.Add(EP.distanceRewardScale, new EnvPara(EnvParaType.FLOAT, 0.0003f));

    //     //The scale of the angularVelocity punish
    //     //angularVelocity punish(every action) = -Tanh(Abs(angularVelocity.sqrMagnitude * 0.1f)) * scale
    //     envParaDict.Add(EP.angularVelocityPunishScale, new EnvPara(EnvParaType.FLOAT, 0.0003f));

    //     //The scale of the time consume punish
    //     //time consume punish(every action) = -scale * 1
    //     envParaDict.Add(EP.timeConsumePunishScale, new EnvPara(EnvParaType.FLOAT, 0.0003f));

    //     //If unifiedEpisode=true, you should using env.reset() to Begin an episode
    //     //In detail, if unifiedEpisode=true, an agent would not immediately reborn after finishing/failing the task. 
    //     //if unifiedEpisode = false, it would.
    //     //如果你用mlagent风格的代码训练，选false
    //     //如果你用gym风格的代码训练，选true
    //     envParaDict.Add(EP.unifiedEpisode, new EnvPara(EnvParaType.BOOL, false));

    //     //Whether to end the episode, if airplane collide with badSphere
    //     envParaDict.Add(EP.isCrashEndEpisode, new EnvPara(EnvParaType.BOOL, true));

    // }

    //easy_fly version settings
    public EnvParaManager()
    {
        eps = Academy.Instance.EnvironmentParameters;

        envParaDict = new Dictionary<EP, EnvPara>();

        envParaDict.Add(EP.seed, new EnvPara(EnvParaType.INT, 5050));
        envParaDict.Add(EP.playgroundNumber, new EnvPara(EnvParaType.INT, 10));

        //*All the position below is localPosition, which being relative to the playground gameObject
        //Target position  
        envParaDict.Add(EP.targetPosX, new EnvPara(EnvParaType.FLOAT, 0f));
        envParaDict.Add(EP.targetPosY, new EnvPara(EnvParaType.FLOAT, 20f));
        envParaDict.Add(EP.targetPosZ, new EnvPara(EnvParaType.FLOAT, 150f));

        //Agent's spawn position 
        envParaDict.Add(EP.spawnPosX, new EnvPara(EnvParaType.FLOAT, 0f));
        envParaDict.Add(EP.spawnPosY, new EnvPara(EnvParaType.FLOAT, 10f));
        envParaDict.Add(EP.spawnPosZ, new EnvPara(EnvParaType.FLOAT, 0f));
        
        //Set random rotation every time agent respawn
        envParaDict.Add(EP.enableRandomAirplaneSpawnRotation, new EnvPara(EnvParaType.BOOL, false));

        //Agents should only fly in a fixed cuboid area.
        //If agents fly out the area, the task failed.
        envParaDict.Add(EP.enableBoundaryCheck, new EnvPara(EnvParaType.BOOL, true));
        envParaDict.Add(EP.boundaryForward, new EnvPara(EnvParaType.FLOAT, 170f));
        envParaDict.Add(EP.boundaryBackward, new EnvPara(EnvParaType.FLOAT, -10f));
        envParaDict.Add(EP.boundaryUp, new EnvPara(EnvParaType.FLOAT, 100f));
        envParaDict.Add(EP.boundaryDown, new EnvPara(EnvParaType.FLOAT, -10f));
        envParaDict.Add(EP.boundaryLeft, new EnvPara(EnvParaType.FLOAT, -40f));
        envParaDict.Add(EP.boundaryRight, new EnvPara(EnvParaType.FLOAT, 40f));

        //If agents is too far from target, the task failed.
        envParaDict.Add(EP.maxDistanceToTarget, new EnvPara(EnvParaType.FLOAT, 1000f));

        //The number of agent per playground, only support 1 recently.
        envParaDict.Add(EP.agentNumberPerPlayground, new EnvPara(EnvParaType.INT, 1));

        //z coordinate of badSpheres. Also indicate the number of badSpheres.ff
        envParaDict.Add(EP.badSpheresPosZ, new EnvPara(EnvParaType.LIST_FLOAT, new List<float>{70,75,80,85}));


        //The x,y coordinates and scale of badSphere is random. 
        envParaDict.Add(EP.badSpherePosX, new EnvPara(EnvParaType.FLOAT, 0f));
        envParaDict.Add(EP.badSpherePosY, new EnvPara(EnvParaType.FLOAT, 10f));
        envParaDict.Add(EP.badSphereScale, new EnvPara(EnvParaType.FLOAT, 25f));

        //The scale of the distance reward
        //distance reward(every action) = Tanh(distanceToTarget * 0.001f) * scale
        envParaDict.Add(EP.distanceRewardScale, new EnvPara(EnvParaType.FLOAT, 0.001f));

        //The scale of the angularVelocity punish
        //angularVelocity punish(every action) = -Tanh(Abs(angularVelocity.sqrMagnitude * 0.1f)) * scale
        envParaDict.Add(EP.angularVelocityPunishScale, new EnvPara(EnvParaType.FLOAT, 0.001f));

        //The scale of the time consume punish
        //time consume punish(every action) = -scale * 1
        envParaDict.Add(EP.timeConsumePunishScale, new EnvPara(EnvParaType.FLOAT, 0.001f));

        //If unifiedEpisode=true, you should using env.reset() to Begin an episode
        //In detail, if unifiedEpisode=true, an agent would not immediately reborn after finishing/failing the task. 
        //if unifiedEpisode = false, it would.
        //如果你用mlagent风格的代码训练，选false
        //如果你用gym风格的代码训练，选true
        envParaDict.Add(EP.unifiedEpisode, new EnvPara(EnvParaType.BOOL, false));

        //Whether to end the episode, if airplane collide with badSphere
        envParaDict.Add(EP.isCrashEndEpisode, new EnvPara(EnvParaType.BOOL, true));

    }

    // **IMPORTATNT: All the Environment Parameters should be got by using GetEnvPara()
    // resmaple only make sense when the parameter is a random parameter(type = RANDOM_UNIFORM or RANDOM_NORMAL)
    // resample=true: the random parameter will be resampled, stored and returned.
    // resample=false: the random parameter will be returned.
    public object GetPara(EP key, bool resample=false)
    {
        if(!envParaDict.ContainsKey(key))
            Debug.LogError("[GetEnvPara]: There is no "+key+" in environment parameters.");
        EnvPara envPara = envParaDict[key];

        //--------------------------------------------------------------------------------------------------
        //--------------------------------------------------------------------------------------------------
        // 由于历史原因，对于浮点数值、整数值、布尔值的参数，此处优先使用MLAgent自带的EnvironmentParameterChannel，
        // 即直接返回Acadamy.Instance.EnvironmentParameters里存储的值。
        // 如果EnvironmentParameterChannel在python端没有设置值，则这里再使用字典里存储的值。
        // 而且由于随机数一定是需要和python端的设置相对应的，所以这样的实现就不能保证EnvParaType.RANDOM的参数
        // 一定是随机数了。如果python那边设置的仅仅是常量浮点数，在C#端是检测不出来的，就会直接返回浮点数。
        
        float para = eps.GetWithDefault(key.ToString(), float.NaN);
        if(!para.Equals(float.NaN))
        {
            if(envPara.type == EnvParaType.FLOAT || envPara.type == EnvParaType.RANDOM_NORMAL || envPara.type == EnvParaType.RANDOM_UNIFORM)
                return (object)para;
            if(envPara.type == EnvParaType.BOOL)
                return (object)(para>0);  
            if(envPara.type == EnvParaType.INT)
                return (object)(int)para;
        }
        //--------------------------------------------------------------------------------------------------
        //--------------------------------------------------------------------------------------------------

        if(resample)
        {
            if(envPara.type != EnvParaType.RANDOM_NORMAL && envPara.type != EnvParaType.RANDOM_UNIFORM)
            {
                //Debug.Log("[GetEnvPara]: The parameter "+key+" is not a random parameter, simply returned the value.");
            }
            else 
            {
                List<float> value = (List<float>)envPara.value;
                if(envPara.type == EnvParaType.RANDOM_NORMAL)
                {
                    value[2] = MathUtils.NormalDisSample(value[0], value[1]);
                }
                else if(envPara.type == EnvParaType.RANDOM_UNIFORM)
                {
                    value[2] = Random.Range(value[0], value[1]);
                }
                return (object)(value[2]);
            }
        }
        if(envPara.type == EnvParaType.RANDOM_NORMAL || envPara.type == EnvParaType.RANDOM_UNIFORM)
            return (object)((List<float>)envPara.value)[2];

        return envPara.value;
    }


    public void SetPara(EP key, EnvParaType ept, object value)
    {
        if(value == null)
            Debug.LogError("[GetEnvPara]: The value is null.");
        if(!envParaDict.ContainsKey(key))
            Debug.LogError("[GetEnvPara]: There is no "+key+" in environment parameters.");
        try
        {
            if(ept == EnvParaType.BOOL)
                if(!(value is bool))
                    throw new UnityException("bool");
            else if(ept == EnvParaType.FLOAT)
                if(!(value is float))
                    throw new UnityException("float");
            else if(ept == EnvParaType.INT)
                if(!(value is int))
                    throw new UnityException("int");
            else if(ept == EnvParaType.LIST_FLOAT)
            {
                List<float> li = value as List<float>;
                if(li == null)
                    throw new UnityException("List<float>");
            }
            else if(ept == EnvParaType.RANDOM_NORMAL || ept == EnvParaType.RANDOM_UNIFORM)
            {
                List<float> li = value as List<float>;
                if(li == null)
                    throw new UnityException("List<float>");
                if(li.Count != 3)
                    throw new UnityException("List<float> (with 3 elements)");
            }
            envParaDict[key].value = value;
            envParaDict[key].type = ept;
        }
        catch(UnityException ue)
        {
            Debug.LogError("[GetEnvPara]: The parameter "+key+"'s value type is mismatched, it should be "+ue.Message+", but it is "+value.GetType().ToString()+".");
        }

    }


}